{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Compress and Evaluate Large Language Models"
   ]
  },
  {
   "cell_type": "raw",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "<a target=\"_blank\" href=\"https://colab.research.google.com/github/PrunaAI/pruna/blob/v|version|/docs/tutorials/llms.ipynb\">\n",
    "    <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
    "</a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "| Component | Details |\n",
    "|-----------|---------|\n",
    "| **Goal** | Show a standard workflow for optimizing and evaluating a large language model |\n",
    "| **Model** | [HuggingFaceTB/SmolLM2-360M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct) |\n",
    "| **Dataset** | [SmolSmolTalk](https://huggingface.co/datasets/HuggingFaceTB/smol-smoltalk) |\n",
    "| **Device** | 1 x RTX A5000 (24GB VRAM) |\n",
    "| **Optimization Algorithms** | quantizer(hqq), compiler(torch_compile) |\n",
    "| **Evaluation Metrics** | perplexity, throughput, total time, energy consumption |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "raw"
    }
   },
   "source": [
    "## Getting Started\n",
    "\n",
    "To install the dependencies, run the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install pruna"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's also set the device to the best available device to make the most out of the optimization process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Load the model\n",
    "\n",
    "Before we can optimize the model, we need to ensure that we can load the model and tokenizer correctly and that they can fit in memory. For this example, we will use a nice and small LLM, [HuggingFaceTB/SmolLM2-360M-Instruct](https://huggingface.co/HuggingFaceTB/SmolLM2-360M-Instruct), but feel free to use any [text-generation model on Hugging Face](https://huggingface.co/models?pipeline_tag=text-generation). \n",
    "\n",
    "Although Pruna works at least as good with much larger models, like Qwen or LLaMA, a small model is a good starting point to show and test the steps of the optimization process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "model_name = \"HuggingFaceTB/SmolLM2-360M-Instruct\"\n",
    "pipe = pipeline(\n",
    "    task=\"text-generation\",\n",
    "    model=model_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we've loaded the model and tokenizer. Let's see if we can run some inference with them. To make this easy for use, we will be using the `transformers` library's `pipeline.__call__` function and passing in a list of messages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "messages = [\n",
    "    {\"role\": \"user\", \"content\": \"Who are you?\"},\n",
    "]\n",
    "pipe(messages, max_new_tokens=100)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the model is able to generate a response to the user's question, which is being cut-off after the allowed `max_new_tokens`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Define the SmashConfig\n",
    "\n",
    "Now we know the model is working, let's continue with the optimization process and define the `SmashConfig`, which we will use later on to optimize the model.\n",
    "\n",
    "Not all optimization algorithms are available for all models but we can learn a bit more about different optimization algorithms and their requirements in the [Algorithms Overview](https://docs.pruna.ai/en/stable/compression.html) section of the documentation.\n",
    "\n",
    "For the current optimization, we will be using the [hqq quantizer](https://docs.pruna.ai/en/stable/compression.html#hqq) and the [torch_compile compiler](https://docs.pruna.ai/en/stable/compression.html#torch-compile). We will updating some parameters for these algorithms, setting `hqq_weight_bits` to `4`, `hqq_compute_dtype` to `torch.bfloat16`, `torch_compile_fullgraph` to `True`, `torch_compile_dynamic` to `True`, and `torch_compile_mode` to `max-autotune`. This is one of the many configurations and will just serve as an example.\n",
    "\n",
    "Let's define the `SmashConfig` object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pruna import SmashConfig\n",
    "\n",
    "smash_config = SmashConfig(\n",
    "    {\n",
    "        \"hqq\": {\"weight_bits\": 4, \"compute_dtype\": \"torch.bfloat16\"},\n",
    "        \"torch_compile\": {\"fullgraph\": True, \"dynamic\": True, \"mode\": \"max-autotune\"}\n",
    "    },\n",
    "    device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Smash the model\n",
    "\n",
    "Now that we have defined the `SmashConfig` object, we can smash the model. We will be using the `smash` function to smash the model and pass the `model` and `smash_config` to it. We also make a deep copy of the model to avoid modifying the original model. \n",
    "\n",
    "Let's smash the model, which should take around 20 seconds for this configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import copy\n",
    "\n",
    "from pruna import smash\n",
    "\n",
    "copy_model = copy.deepcopy(pipe.model).to(\"cpu\")\n",
    "smashed_model = smash(\n",
    "    model=pipe.model,\n",
    "    smash_config=smash_config,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we've optimized the model. Let's see if everything still works as expected and we can run some inference with the optimized model. In this case, we are running the inference by first encoding the prompt through the `tokenizer` and then passing the `input_ids` to the `PrunaModel.generate` method, which also allows us to specify additional parameters such as `max_new_tokens`.\n",
    "\n",
    "If you are using `torch_compile` as your compiler, you can expect the first inference warmup to take a bit longer than the actual inference."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Who are you?\"\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "pipe(messages, max_new_tokens=256)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the model is able to generate a similar response to the original model. \n",
    "\n",
    "If you notice a significant difference, it might have several reasons, the models, the configuration, the hardware, etc. As optimization can be non-deterministic, we encourage you to retry the optimization process or try out different configurations and models to find the best fit for your use case but also feel free to reach out to us on [Discord](https://discord.gg/JFQmtFKCjd) if you have any questions or feedback."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Evaluate the smashed model\n",
    "\n",
    "Now that we have optimized the model, we can evaluate the performance of the optimized model. We will be using the `EvaluationAgent` to evaluate the performance of the optimized model. We will do so with some basic metrics, the `elapsed_time`, as well as a stateful metrics, the `perplexity`. An overview of the different metrics can be found in our [documentation](https://docs.pruna.ai/).\n",
    "\n",
    "Let's define the `EvaluationAgent` object and start the evaluation process. Note that we are using the `datamodule.limit_datasets(100)` method to limit the number of datasets to 100, which is just for the sake of time. Additionally, set the `n_iterations` and `n_warmup_iterations` to ensure that we monitor the performance of the model whenever it is running smoothly.\n",
    "\n",
    "The evaluation can take anywhere from a couple of minutes to a couple of hours to complete, depending on your hardware, the number of samples in the dataset, and the configuration of the model. In our case it should only take a couple of minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pruna import PrunaModel\n",
    "from pruna.data.pruna_datamodule import PrunaDataModule\n",
    "from pruna.engine.utils import move_to_device\n",
    "from pruna.evaluation.evaluation_agent import EvaluationAgent\n",
    "from pruna.evaluation.metrics import (\n",
    "    EnergyConsumedMetric,\n",
    "    ThroughputMetric,\n",
    "    TorchMetricWrapper,\n",
    "    TotalTimeMetric,\n",
    ")\n",
    "from pruna.evaluation.task import Task\n",
    "\n",
    "# Define the metrics\n",
    "metrics = [\n",
    "    EnergyConsumedMetric(n_iterations=50, n_warmup_iterations=5),\n",
    "    ThroughputMetric(n_iterations=50, n_warmup_iterations=5),\n",
    "    TotalTimeMetric(n_iterations=50, n_warmup_iterations=5),\n",
    "    TorchMetricWrapper(\"perplexity\", call_type=\"single\"),\n",
    "]\n",
    "\n",
    "# Define the datamodule\n",
    "pipe.tokenizer.pad_token = pipe.tokenizer.eos_token\n",
    "datamodule = PrunaDataModule.from_string(\"SmolSmolTalk\", tokenizer=pipe.tokenizer)\n",
    "datamodule.limit_datasets(100)\n",
    "\n",
    "# Define the task and evaluation agent\n",
    "task = Task(metrics, datamodule=datamodule, device=device)\n",
    "eval_agent = EvaluationAgent(task)\n",
    "\n",
    "# Update the model args to ensure right generation arguments are passed after compilation\n",
    "inference_args = {\"max_new_tokens\": 250}\n",
    "\n",
    "# Evaluate smashed model and offload it to CPU\n",
    "move_to_device(smashed_model, device)\n",
    "smashed_model.inference_handler.model_args.update(inference_args)\n",
    "smashed_model_results = eval_agent.evaluate(smashed_model)\n",
    "move_to_device(smashed_model, \"cpu\")\n",
    "\n",
    "# Evaluate base model and offload it to CPU\n",
    "base_model = PrunaModel(model=copy_model)\n",
    "move_to_device(base_model, device)\n",
    "base_model.inference_handler.model_args.update(inference_args)\n",
    "base_model_results = eval_agent.evaluate(base_model)\n",
    "move_to_device(base_model, \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can see the results of the evaluation and compare the performance of the original and the optimized model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Markdown, display  # noqa\n",
    "\n",
    "\n",
    "def make_comparison_table(base_model_results, smashed_model_results):  # noqa\n",
    "    header = \"| Metric | Base Model | Smashed Model | Improvement % |\\n\"\n",
    "    header += \"|\" + \"-----|\" * 4 + \"\\n\"\n",
    "    rows = []\n",
    "\n",
    "    for base, smashed in zip(base_model_results, smashed_model_results):\n",
    "        base_result = base.result\n",
    "        smashed_result = smashed.result\n",
    "        if base.higher_is_better:\n",
    "            diff = ((smashed_result - base_result) / base_result) * 100\n",
    "        else:\n",
    "            diff = ((base_result - smashed_result) / base_result) * 100\n",
    "        row = f\"| {base.name} | {base_result:.4f} {base.metric_units or ''}\"\n",
    "        row += f\"| {smashed_result:.4f} {smashed.metric_units or ''} | {diff:.2f}% |\"\n",
    "        rows.append(row)\n",
    "    return header + \"\\n\".join(rows)\n",
    "\n",
    "\n",
    "display(Markdown(make_comparison_table(base_model_results, smashed_model_results)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As we can see, the optimized model handles 4x more throughput and consumes only 1/5 of the energy of the base model, while losing only a small portion of its performance based on the perplexity metric, which is expected given the nature of the optimization. Now, we can start to compare, iterate and see what optimization works best for our models, given the metrics we are interested in.\n",
    "\n",
    "We can now save the optimized model to disk and share it with others. Note that some optimizations, such as `torch_compile`, are device dependent and will be re-applied when loading the model on a different device."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the model either locally or push it to huggingface hub\n",
    "smashed_model.save_pretrained(\"smashed_model\")\n",
    "# smashed_model.push_to_hub(\"PrunaAI/smashed_model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wrap up\n",
    "\n",
    "In this tutorial, we have shown a standard workflow for optimizing and evaluating a large language model. We have used the `SmashConfig` object to define the optimization algorithms and the `EvaluationAgent` to evaluate the performance of the optimized model. We have also used the `PrunaDataModule` to load the dataset and the `Task` object to define the task and evaluation agent.\n",
    "\n",
    "We have shown how to optimize the model using the `smash` function and how to evaluate the performance of the optimized model using the `EvaluationAgent`.\n",
    "\n",
    "Proving we can optimize the model, by making it quicker, more energy efficient and using less memory, while only losing a small amount of accuracy."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pruna",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
